import pyudev

ALL_DEVICE_INFO = 0x1

CLASS_HD = 0x10
CLASS_ATA = 0x100
CLASS_USB = 0x1000
CLASS_OTH = 0x10000

CLASS_CDROM = 0x100000
CLASS_FLOPPY = 0x1000000


def get_all_devices():
    context = pyudev.Context()
    udev_devices = context.list_devices()
    devices_info = {}
    devices_size = {}
    for device in udev_devices:
        dev_info = {}
        udev_name = device.keys()
        for n in udev_name:
            try:
                dev_info[n] = device[n]
            except:
                pass
        devices_info[device.sys_name] = dev_info
        try:
             devices_size[device.sys_name] = int(device.attributes['size']) * 512
        except:
            pass
    return devices_info, devices_size

def probe(dev_class):
    devices, devices_size = get_all_devices()

    if dev_class & ALL_DEVICE_INFO:
        return devices, devices_size

    if dev_class & CLASS_HD or \
            dev_class & CLASS_ATA or \
            dev_class & CLASS_USB:
        ata_hd_map, usb_hd_map, oth_hd_map = get_hd_ata_usb(devices)
        hd_map = {}
        hd_map.update(ata_hd_map)
        hd_map.update(usb_hd_map)
        hd_map.update(oth_hd_map)
        if dev_class & CLASS_HD:
            return hd_map
        if dev_class & CLASS_ATA:
            return ata_hd_map
        if dev_class & CLASS_USB:
            return usb_hd_map
        if dev_class & CLASS_OTH:
            return oth_hd_map
    if dev_class & CLASS_CDROM:
        return get_cdrom(devices)
    if dev_class & CLASS_FLOPPY:
        return get_floppy(devices)

def get_floppy(devices):
    # Because I have no floppy, so I can not test it...
    floppy_map = {}
    for devname in devices.keys():
        device = devices[devname]
        if 'ID_TYPE' in device.keys() and device['ID_TYPE'] == 'floppy':
            floppy_map[devname] = []
    return floppy_map

def get_cdrom(devices):
    cdrom_map = {}
    for devname in devices.keys():
        device = devices[devname]
        if 'ID_TYPE' in device.keys() and device['ID_TYPE'] == 'cd':
            cdrom_map[devname] = []
    return cdrom_map
        
def get_hd_ata_usb(devices):
    part_info_mm = {}
    disk_info_mm = {}
    disk_size_map = {}
    for devname in devices.keys():
        device = devices[devname]
        if 1:
            if 'ID_TYPE' in device.keys() and device['ID_TYPE'] == 'disk':
                if device['DEVTYPE'] == 'partition':
                    part_info_mm[devname] = device
                elif device['DEVTYPE'] == 'disk':
                    disk_info_mm[devname] = device

    usb_hd_map = {}
    ata_hd_map = {}
    oth_hd_map = {}
    for key in disk_info_mm.keys():
        for k in disk_info_mm[key].keys():
            if k == 'ID_BUS':
                if disk_info_mm[key][k] == 'usb':
                    usb_hd_map[key] = []
                elif disk_info_mm[key][k] == 'ata':
                    ata_hd_map[key] = []
                else:
                    oth_hd_map[key] = []

    for key in part_info_mm.keys():
        devpath = part_info_mm[key]['DEVPATH']
        disk_name = 'None'
        try:
            disk_name = devpath[:devpath.rfind('/')]
            disk_name = disk_name[disk_name.rfind('/')+1:]
        except:
            pass
        if disk_name:
            if usb_hd_map.has_key(disk_name):
                usb_hd_map[disk_name].append(key)
            elif ata_hd_map.has_key(disk_name):
                ata_hd_map[disk_name].append(key)
            elif oth_hd_map.has_key(disk_name):
                oth_hd_map[disk_name].append(key)
    return ata_hd_map, usb_hd_map, oth_hd_map



if __name__ == '__main__':
    print 'The hd information:'
    print probe(CLASS_HD)

    print 'The ata information:'
    print probe(CLASS_ATA)

    print 'The usb devices:'
    print probe(CLASS_USB)

    print 'The cdrom devices:'
    print probe(CLASS_CDROM)
    
    print 'The floppy devices:'
    print probe(CLASS_FLOPPY)
    
    print 'Disk Size:'
    devices_info, devices_size = probe(ALL_DEVICE_INFO)#get_all_devices()
    print devices_size
